#ifndef _h_tree
#define _h_tree

#include "geo.h"
#include <vector>
#include <algorithm>
using namespace std;

const size_t TREE_LEAF_LIM=10;
const double TREE_SAH_KT=1.;
const double TREE_SAH_KI=100.;
const int TREE_MAX_DEPTH=20;

class TreeNode {
	TreeNode *lc, *rc;
	vector<Patch> patches;
	Box box;
	int cut_dim; double cut_posi;
	
	struct EventPoint {
		int type;
		double posi;
		EventPoint(int t, double p)
			: type(t), posi(p) {}
		bool operator <(const EventPoint &e) const {
			return posi<e.posi;
		}
	};
	
	double SAH_C(double pl, double pr, int nl, int nr) {
		const double KT=TREE_SAH_KT;
		const double KI=TREE_SAH_KI;
		double c=KT+KI*(pl*nl+pr*nr);
		//if(nl==0 || nr==0) c*=0.8;
		return c;
	}
	
	double SAH(int dim, double posi, Box box, int nl, int nr, int np) {
		double area=box.get_surface_area();
		Box bl, br;
		if(!box.split(dim, posi, bl, br))
			return INFTY;
		double pl=bl.get_surface_area()/area;
		double pr=br.get_surface_area()/area;
		double c=SAH_C(pl, pr, nl+np, nr);
		return c;
	}
public:
	TreeNode(vector<Patch> &patches, Box box, int depth) {
		//printf("TreeNode: building (patches.size=%d)\n", patches.size());
		this->box=box;
		// find the best cut
		double min_cost=INFTY;
		int min_dim=-1;
		double min_posi=0.;
		for(int dim=0; dim<3; ++dim) {
			vector<EventPoint> events;
			for(size_t i=0; i<patches.size(); ++i) {
				pair<double, double> bound=patches[i].bound(dim);
				if(bound.first==bound.second)
					events.push_back(EventPoint(2, bound.first));
				else {
					events.push_back(EventPoint(0, bound.first));
					events.push_back(EventPoint(1, bound.second));
				}
			}
			sort(events.begin(), events.end());
			int nl=0, nr=(int)patches.size();
			for(size_t i=0; i<events.size(); ) {
				int nplus=0, nminus=0, np=0;
				double posi=events[i].posi;
				while(i<events.size() && events[i].posi==posi) {
					if(events[i].type==0) ++nplus;
					if(events[i].type==1) ++nminus;
					if(events[i].type==2) ++np;
					++i;
				}
				nr-=nminus+np;
				double cost=SAH(dim, posi, box, nl, nr, np);
				nl+=nplus+np;
				if(cost<min_cost) {
					min_cost=cost;
					min_posi=posi;
					min_dim=dim;
				}
			}
		}
		// boundary condition
		if(patches.size()<TREE_LEAF_LIM || min_cost>patches.size()*TREE_SAH_KI || depth>TREE_MAX_DEPTH) {
			lc=rc=NULL;
			this->patches=patches;
			return;
		}
		// okay now we can cut it
		cut_dim=min_dim;
		cut_posi=min_posi;
		vector<Patch> l_patches;
		vector<Patch> r_patches;
		for(size_t i=0; i<patches.size(); ++i) {
			pair<double, double> bound=patches[i].bound(min_dim);
			if(bound.first==bound.second) {
				if(bound.first<=min_posi)
					l_patches.push_back(patches[i]);
				else
					r_patches.push_back(patches[i]);
			} else {
				if(bound.first<min_posi)
					l_patches.push_back(patches[i]);
				if(bound.second>min_posi)
					r_patches.push_back(patches[i]);					
			}
		}
		patches.clear(); // free the memory
		Box l_box, r_box;
		if(!box.split(min_dim, min_posi, l_box, r_box))
			assert(false);
		lc=new TreeNode(l_patches, l_box, depth+1);
		rc=new TreeNode(r_patches, r_box, depth+1);		
	}
	~TreeNode() {
		if(lc!=NULL) delete lc;
		if(rc!=NULL) delete rc;
	}
	bool get_cross(Ray ray, Patch *patch, double *t, double *alpha, double *beta) const {
		if(!box.check_ray(ray))
			return false;
		if(lc==NULL) { // is leaf
			bool found=false;
			double ct, ca, cb;
			for(size_t i=0; i<patches.size(); ++i)
				if(patches[i].get_cross(ray, &ct, &ca, &cb)) {
					if(!box.inside(patches[i].get_point(ca, cb)))
						continue;
					if(!found || ct<*t) {
						*t=ct; *alpha=ca, *beta=cb;
						*patch=patches[i];
						found=true;
					}
				}
			return found;
		}
		if(ray.s.get_dim(cut_dim)<=cut_posi) {
			if(lc->get_cross(ray, patch, t, alpha, beta))
				return true;
			return rc->get_cross(ray, patch, t, alpha, beta);
		} else {
			if(rc->get_cross(ray, patch, t, alpha, beta))
				return true;
			return lc->get_cross(ray, patch, t, alpha, beta);
		}
	}
	bool check_cross(Ray ray) const {
		if(!box.check_ray(ray))
			return false;
		if(lc==NULL) { // is leaf
			for(size_t i=0; i<patches.size(); ++i)
				if(patches[i].check_cross(ray))
					return true;
			return false;
		}
		if(lc->check_cross(ray))
			return true;
		if(rc->check_cross(ray))
			return true;
		return false;
	}
};

class Tree {
	//vector<Patch> patches;
	TreeNode *root;
public:
	Tree(const vector<Patch> &patches) {
		printf("Tree: start build\n");
		Box box;
		for(size_t i=0; i<patches.size(); ++i)
			for(int j=0; j<3; ++j)
				box.add(patches[i].p[j]);
		vector<Patch> p=patches;
		root=new TreeNode(p, box, 1);
		printf("Tree: end build\n");
	}
	~Tree() {
		delete root;
	}
	bool get_cross(Ray ray, Patch *patch, double *t, double *alpha, double *beta) const {
		return root->get_cross(ray, patch, t, alpha, beta);
		/*
		bool found=false;
		for(size_t i=0; i<patches.size(); ++i) {
			Point cp; double ct;
			if(!patches[i].get_cross(ray, &cp, &ct))
				continue;
			if(!found || ct<*t) {
				*t=ct;
				*p=cp;
				*patch=patches[i];
				found=true;
			}
		}
		return found;
		*/
	}
	bool check_cross(Ray ray) const {
		return root->check_cross(ray);
	}
};

#endif
